import os
import uuid
import json
import time
import jwt
import requests
import cv2
from typing import Optional, Tuple, Dict, Any

import urllib3
from dotenv import load_dotenv
from openai import OpenAI
from utils import sr_s3
from utils.sr_utils import image_path_to_data_url, video_compress_image_size_path, image_examine, content_moderations


class LumaVideoGenerator:
    """视频生成工具类，封装Luma Labs视频生成API调用逻辑，支持外部调用"""

    def __init__(self):
        """初始化配置与环境变量"""
        # 加载环境变量
        self.verify_ssl = True
        load_dotenv(verbose=True, override=True)
        self.app_env = os.getenv('APP_ENV', 'prod')
        self.url = "https://api.lumalabs.ai"  # Luma Labs API基础地址
        self.luma_api_key = os.getenv('LUMA_API_KEY')  # Luma API密钥
        self.NSFW_CHECK_ENABLED = os.getenv('NSFW_CHECK_ENABLED')

        # 初始化鉴权头（Luma使用Bearer认证）
        self.headers = {
            "Authorization": f"Bearer {self.luma_api_key}",
            "Content-Type": "application/json"
        }

        # 临时文件存储目录
        self.temp_image_dir = "./upload"
        self.temp_video_dir = "./download"
        os.makedirs(self.temp_image_dir, exist_ok=True)
        os.makedirs(self.temp_video_dir, exist_ok=True)

    # --------------------------
    # 内部工具方法：文件处理
    # --------------------------
    def _delete_temp_files(self, *file_paths: str) -> None:
        """删除临时文件（忽略不存在的文件）"""
        for path in file_paths:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"清理临时文件失败 {path}: {str(e)}")

    def _download_and_compress_image(self, img_url: str, storage_service: str, bucket_name: str) -> Tuple[bool, str, int]:
        """下载图片并压缩（返回是否成功、图片URL、状态码）"""
        from utils import sr_s3

        if self.NSFW_CHECK_ENABLED == "true":
            image_url = sr_s3.get_url(img_url, serviceName=storage_service, bucketName=bucket_name)
            clientOpenAI = OpenAI()
            image_result = content_moderations(clientOpenAI, image_url)
            if image_result:
                return False, f"图片违规", 5
        try:
            # 生成临时保存路径
            name = str(uuid.uuid4())
            ext = os.path.splitext(img_url)[-1].lower() or '.jpg'
            save_path = os.path.join(self.temp_image_dir, f"{name}{ext}")
            # 下载图片
            success, err = sr_s3.download_file(
                img_url, save_path, serviceName=storage_service, bucketName=bucket_name
            )
            if not success:
                return False, f"下载失败: {err}", 2
            # 压缩图片
            success, new_path = video_compress_image_size_path(save_path, 10, (300, 6000), (300, 6000))
            if not success:
                return False, "图片压缩失败", 2
            # 上传压缩后的图片获取可访问URL（Luma支持URL输入）
            from utils import sr_s3
            s3_key = sr_s3.upload_login_or_visitor(
                new_path, "temp/images", serviceName=storage_service, bucketName=bucket_name
            )
            image_url = sr_s3.get_url(s3_key, serviceName=storage_service, bucketName=bucket_name)
            # self._delete_temp_files(new_path, save_path)
            return True, image_url, 1
        except Exception as e:
            return False, f"图片处理异常: {str(e)}", 2

    def _download_video(self, video_url: str, task_id: str) -> Tuple[bool, Dict[str, Any]]:
        """下载生成的视频到本地"""
        try:
            filename = f"video_{task_id}.mp4"
            save_path = os.path.join(self.temp_video_dir, filename)

            response = requests.get(video_url, stream=True, timeout=30)
            response.raise_for_status()

            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        f.write(chunk)

            # 获取视频信息
            cap = cv2.VideoCapture(save_path)
            video_info = {
                "save_path": save_path,
                "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "fps": cap.get(cv2.CAP_PROP_FPS),
                "duration": cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0
            }
            cap.release()
            return True, video_info
        except Exception as e:
            return False, {"error": f"视频下载失败: {str(e)}"}

    def _upload_video_to_s3(self, video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        """上传视频到S3"""
        if not video_path or not os.path.exists(video_path):
            return False, "视频文件不存在"
        try:
            s3_key = sr_s3.upload_login_or_visitor(
                video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
            )
            return True, s3_key
        except Exception as e:
            return False, f"上传S3失败: {str(e)}"

    # --------------------------
    # 核心方法：API调用与任务处理
    # --------------------------
    def _create_luma_video_task(self, payload: dict) -> Tuple[bool, str, str]:
        """调用Luma Labs API创建视频任务"""
        try:
            url = f"{self.url}/dream-machine/v1/generations/video"  # Luma创建视频端点
            print(url)
            print(payload)
            # exit()
            response = requests.post(url, headers=self.headers, data=json.dumps(payload))
            response_data = response.json()
            print(response_data)
            # 处理Luma API响应（参考文档：成功返回带id的对象）
            if response.status_code == 201:
                task_id = response_data.get("id")
                if task_id:
                    return True, task_id, ""
                else:
                    return False, "", "API返回成功但无任务ID"
            else:
                error_msg = response_data.get("error", {}).get("message", "未知错误")
                return False, "", f"创建任务失败: {response_data} {error_msg} (状态码: {response.status_code})"
        except Exception as e:
            return False, "", f"创建任务异常: {str(e)}"

    def _poll_task_status(self, task_id: str, timeout: int = 600, interval: int = 5) -> Tuple[bool, str]:
        """轮询Luma任务状态"""
        if not task_id:
            return False, "任务ID为空"

        url = f"{self.url}/dream-machine/v1/generations/{task_id}"  # Luma查询任务端点
        start_time = time.time()
        retries = 0
        max_retries = 5

        session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            max_retries=max_retries,
            pool_connections=10,
            pool_maxsize=10
        )
        session.mount('https://', adapter)

        if not self.verify_ssl:
            urllib3.disable_warnings()

        while time.time() - start_time < timeout:
            try:
                response = session.get(
                    url,
                    headers=self.headers,
                    timeout=(5, 15),
                    verify=self.verify_ssl
                )
                response.raise_for_status()
                response_data = response.json()
                print(response_data)
                # Luma状态可能为：processing, completed, failed, cancelled
                status = response_data.get("state")
                print(status)
                if status == "completed":
                    # 成功时视频URL在assets->video->url
                    video_assets = response_data.get("assets", {}).get("video")
                    if video_assets:
                        return True, video_assets
                    else:
                        return False, "任务成功但未找到视频URL"
                elif status == "failed":
                    error_msg = response_data.get("failure_reason")
                    return False, error_msg
                elif status in ["cancelled", "expired"]:
                    return False, f"任务状态: {status}"
                else:
                    # 处理中状态：processing, pending等
                    print(f"任务状态: {status}，等待{interval}秒后重试")
                    time.sleep(interval)
                    retries = 0

            except requests.exceptions.RequestException as e:
                retries += 1
                if retries >= max_retries:
                    return False, f"请求失败（重试{retries}次）: {str(e)}"
                print(f"请求异常（{retries}/{max_retries}）: {str(e)}，{interval}秒后重试")
                time.sleep(interval * retries)
            except Exception as e:
                return False, f"轮询异常: {str(e)}"

        return False, f"任务超时（超过{timeout}秒）"

    # --------------------------
    # 外部调用方法：主入口
    # --------------------------
    def process_video_task(self, task_params: Dict[str, Any]):
        """处理视频生成任务（Luma Labs版）"""
        video_info = None
        try:
            # 参数校验
            required_params = ['prompt', 'modelType', 'storageService', 'bucketName', 'storedPrefix', 'QueueResult']
            for param in required_params:
                if not task_params.get(param):
                    return {"status": 0, "msg": f"缺少参数: {param}", 'code': 6003}

            model_type = task_params.get('modelType')
            ratio = task_params.get('ratio')
            payload = {
                "model": task_params['moduleName'],
                "prompt": task_params.get('prompt'),
                "loop": task_params.get('loop', False),
            }
            if task_params['moduleName'] != 'ray-1-6':
                payload['duration'] = task_params.get('duration', '5') + 's',

            if ratio is not None:
                payload['aspect_ratio'] = ratio
            storage_service = task_params.get('storage_service')
            bucket_name = task_params.get('bucketName')
            # 根据模型类型补充参数（Luma支持文生视频和图生视频）
            if model_type == 1:  # 文生视频
                payload['aspect_ratio'] = task_params.get('ratio', "16:9")
            elif model_type in [2, 3]:  # 图生视频（单图/首尾图）
                first_image_url = task_params.get('first_image')
                if not first_image_url:
                    return {"status": 0, "msg": "图生视频需提供first_image", 'code': 6003}
                # 获取处理后的图片URL（Luma需要可公开访问的URL）
                image_url = sr_s3.get_url(first_image_url, serviceName=storage_service, bucketName=bucket_name)
                # 初始化关键帧字典
                payload['keyframes'] = {
                    "frame0": {
                        "type": "image",
                        "url": image_url
                    }
                }
                # 首尾图模式可增加过渡参数（如果Luma支持）
                if model_type == 3 and task_params.get('last_image'):
                    image_url2 = sr_s3.get_url(task_params.get('last_image'), serviceName=storage_service, bucketName=bucket_name)
                    payload['keyframes']['frame1'] = {
                        "type": "image",
                        "url": image_url2
                    }
            else:
                return {"status": 0, "msg": f"不支持的模型类型: {model_type}", 'code': 6003}

            # 创建任务
            print('创建任务##')
            success, task_id, err_msg = self._create_luma_video_task(payload)
            if not success:
                return {"status": 0, "msg": err_msg}
            # task_id = '32c94ebc-fd3f-4acd-875a-aac55782f72e'
            print(f"r任务id:::{task_id}")
            # 轮询结果
            success, video_url = self._poll_task_status(
                task_id,
                timeout=task_params.get('timeout', 600),
                interval=task_params.get('interval', 5)
            )
            if not success:
                if video_url == "400: image moderation failed" or video_url == '400: prompt not allowed because advanced moderation failed':
                    return {"status": 7, "msg": video_url, "code": 7202}
                return {"status": 0, "msg": video_url, "code": 6001}

            # 下载并上传到S3
            success, video_info = self._download_video(video_url, task_id)
            if not success:
                return {"status": 0, "msg": video_info.get("error"), "code": 6001}

            success, s3_key = self._upload_video_to_s3(
                video_info["save_path"],
                task_params['storedPrefix'],
                task_params['storageService'],
                task_params['bucketName']
            )
            if not success:
                return {"status": 0, "msg": s3_key, "code": 6001}

            # 发送结果到队列
            from utils import sr_sqs
            msg_body = {
                'status': 1,
                'width': video_info['width'],
                'height': video_info['height'],
                'url': s3_key,
                'local_path': video_info['save_path'],
                "model_mode": "std",
                'msg': "success",
                'code': ''

            }
            return msg_body

        except Exception as e:
            return {"status": 0, "msg": f"任务处理异常: {str(e)}", "code": 6001}


# 调用示例
if __name__ == "__main__":
    video_generator = LumaVideoGenerator()
    # task_params = {
    #     "prompt": "在看书",
    #     "modelType": 2,  # 1=文生视频, 2=单图生视频,
    #     "modelName": "ray-2",
    #     "duration": '9',
    #     "resolution": "1080",
    #     "aspect_ratio": "16:9",
    #     "storageService": "r2",
    #     "bucketName": "picwand",
    #     "loop": True,
    #     "storedPrefix": "video/visitor/luma",
    #     "QueueResult": "AILocalVideoResult",
    #     # 图生视频时添加
    #     # "first_image": "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754018727525.png",
    #     "first_image": "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1753673942065.png",  # 首帧图片URL
    #
    # }
    task_params = {
        'taskId':'44f7065ba31eaf292b8e',
    'prompt':'街道中出现了两个小孩在玩皮球',
    'first_image':'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1755747499445.jpg',
    'last_image':'',
    'duration':'9',
    'fps':'30',
    'resolution':'720p',
    'ratio':'9:16',
    'appId':'1',
    'taskType':'pwai_video_itf',
    'QueueResult':'AILocalVideoResult',
    'storageService':'r2',
    'isLogin':1,
    'bucketName':'picwand',
    'storedPrefix':'video/image_to_video',
    'progress':1,
    'queueName':'AIHubVideoTasks_Test',
    'modelType':2,
    'moduleType':26,
    'moduleName':'ray-flash-2',
    'mResolution':'540p',
    'mode':'std',
    'negative_prompt':'',
    'cfg_scale':0.5,
    'p_t_b':0,
    'n_t_b':0,
    'motion_type':'',
    'model_style':'',
    'camera_movement':'',
    }
    result = video_generator.process_video_task(task_params)
    print(result)
